package com.liru.chat_demo.service;

import com.liru.chat_demo.model.FunctionSchema;
import com.liru.chat_demo.model.McpServer;
import com.liru.chat_demo.model.McpTool;
import com.liru.chat_demo.repository.McpServerRepository;
import com.liru.chat_demo.repository.McpToolRepository;
import io.modelcontextprotocol.client.McpAsyncClient;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
import io.modelcontextprotocol.spec.McpSchema;
import jakarta.annotation.PostConstruct;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * 用于管理 MCP 客户端和工具的服务
 */
@Service
public class McpClientService {
    // 日志记录器
    private static final Logger logger = LoggerFactory.getLogger(McpClientService.class);
    // 超时时间设置为30秒
    private static final Duration TIMEOUT = Duration.ofSeconds(30);

    // MCP服务器数据访问仓库
    private final McpServerRepository serverRepository;
    // MCP工具数据访问仓库
    private final McpToolRepository toolRepository;
    // 存储客户端的并发安全映射表
    private final Map<String, McpAsyncClient> clientMap = new ConcurrentHashMap<>();

    // 构造函数注入依赖
    public McpClientService(McpServerRepository serverRepository, McpToolRepository toolRepository) {
        this.serverRepository = serverRepository;
        this.toolRepository = toolRepository;
    }

    /**
     * 在应用程序启动时初始化 MCP 客户端
     */
    @PostConstruct
    public void init() {
        // 添加一个默认的 MCP 服务器用于测试
        McpServer defaultServer = new McpServer(
                "默认 MCP 服务器",
                "http://localhost:8999",
                "/api/v1/sse",
                true
        );
        serverRepository.save(defaultServer);
        
        // 为所有启用的服务器初始化客户端
        refreshAllClients();
    }

    /**
     * 刷新所有 MCP 客户端
     */
    public void refreshAllClients() {
        // 获取所有启用的服务器
        Collection<McpServer> enabledServers = serverRepository.findAllEnabled();
        // 遍历每个服务器并刷新其客户端
        for (McpServer server : enabledServers) {
            refreshClient(server.getId());
        }
    }

    /**
     * 根据服务器ID刷新 MCP 客户端
     * @param serverId 服务器ID
     * @return 成功返回true，失败返回false
     */
    public boolean refreshClient(String serverId) {
        // 根据ID查找服务器，如果不存在或未启用则关闭客户端并返回false
        Optional<McpServer> serverOpt = serverRepository.findById(serverId);
        if (serverOpt.isEmpty() || !serverOpt.get().isEnabled()) {
            closeClient(serverId);
            return false;
        }

        McpServer server = serverOpt.get();
        try {
            // 如果存在现有客户端则关闭它
            closeClient(serverId);
            
            // 创建新的客户端传输层
            HttpClientSseClientTransport transport = HttpClientSseClientTransport.builder(server.getBaseUrl())
                    .sseEndpoint(server.getSseEndpoint())
                    .build();
            
            // 构建异步客户端
            McpAsyncClient client = McpClient.async(transport).build();
            
            // 初始化客户端
            client.initialize().block(TIMEOUT);
            
            // 将客户端存储到映射表中
            clientMap.put(serverId, client);
            
            // 刷新工具列表
            refreshTools(serverId);
            
            // 更新最后访问时间并保存服务器信息
            server.updateLastSeen();
            serverRepository.save(server);
            
            return true;
        } catch (Exception e) {
            // 记录错误日志并返回false
            logger.error("刷新服务器 {} 的 MCP 客户端失败: {}", serverId, e.getMessage());
            return false;
        }
    }

    /**
     * 为 MCP 服务器刷新工具列表
     * @param serverId 服务器ID
     */
    private void refreshTools(String serverId) {
        // 获取对应服务器的客户端
        McpAsyncClient client = clientMap.get(serverId);
        // 如果客户端不存在则直接返回
        if (client == null) {
            return;
        }

        try {
            // 删除该服务器现有的所有工具
            toolRepository.deleteAllByServerId(serverId);
            
            // 从服务器获取工具列表
            McpSchema.ListToolsResult toolsResult = client.listTools().block(TIMEOUT);
            // 如果没有获取到工具则记录警告日志并返回
            if (toolsResult == null || toolsResult.tools() == null) {
                logger.warn("服务器 {} 未找到任何工具", serverId);
                return;
            }
            
            // 注册获取到的每个工具
            for (McpSchema.Tool tool : toolsResult.tools()) {
                McpTool mcpTool = new McpTool(
                        serverId,
                        tool.name(),
                        tool.name(),
                        tool.description(),
                        tool.inputSchema() != null ? convertParametersToMap(tool.inputSchema()) : Map.of()
                );
                toolRepository.save(mcpTool);
                logger.info("已注册来自服务器 {} 的工具 {}", serverId, tool.name());
            }
        } catch (Exception e) {
            // 记录工具刷新失败的错误日志
            logger.error("刷新服务器 {} 的工具失败: {}", serverId, e.getMessage());
        }
    }

    /**
     * 将 MCP 参数转换为映射表
     * @param parameters MCP 参数
     * @return 转换后的参数映射表
     */
    private Map<String, Object> convertParametersToMap(McpSchema.JsonSchema parameters) {
        // 创建结果映射表
        Map<String, Object> result = new HashMap<>();
        // 如果参数为空则直接返回空结果
        if (parameters == null) {
            return result;
        }

        // 创建属性映射表和必需参数列表
        Map<String, Object> properties = new HashMap<>();
        List<String> required = new ArrayList<>();

        // 遍历每个参数
        for (Map.Entry<String, Object> entry : parameters.properties().entrySet()) {
            String paramName = entry.getKey();
            LinkedHashMap paramDetail = (LinkedHashMap) entry.getValue();

            // 创建参数模式映射表
            Map<String, Object> paramSchema = new HashMap<>();
            paramSchema.put("type", paramDetail.get("type"));
            paramSchema.put("format", paramDetail.get("format"));
            // 将参数模式添加到属性映射表中
            properties.put(paramName, paramSchema);
        }

        // 设置结果映射表的基本结构
        result.put("type", "object");
        result.put("properties", properties);
        // 如果有必需参数则添加到结果中
        if (!required.isEmpty()) {
            result.put("required", required);
        }
        
        return result;
    }

    /**
     * 根据服务器ID关闭 MCP 客户端
     * @param serverId 服务器ID
     */
    public void closeClient(String serverId) {
        // 从映射表中移除并获取客户端
        McpAsyncClient client = clientMap.remove(serverId);
        // 如果客户端存在则尝试优雅关闭
        if (client != null) {
            try {
                client.closeGracefully().block(TIMEOUT);
                logger.info("已关闭服务器 {} 的 MCP 客户端", serverId);
            } catch (Exception e) {
                // 记录关闭失败的错误日志
                logger.error("关闭服务器 {} 的 MCP 客户端失败: {}", serverId, e.getMessage());
            }
        }
    }

    /**
     * 调用 MCP 工具
     *
     * @param toolName  工具名称
     * @param arguments 工具参数
     * @return 工具执行结果
     */
    public String callTool(String toolName, Map<String, Object> arguments) {
        // 查找指定名称且已暴露的工具
        Optional<McpTool> toolOpt = toolRepository.findAll().stream()
                .filter(tool -> tool.getName().equals(toolName) && tool.isExposed())
                .findFirst();
                
        // 如果未找到工具则返回错误
        if (toolOpt.isEmpty()) {
            return "未找到工具或工具未暴露: " + toolName;
        }
        
        McpTool tool = toolOpt.get();
        // 获取工具对应的客户端
        McpAsyncClient client = clientMap.get(tool.getServerId());
        
        // 如果客户端不可用则返回错误
        if (client == null) {
            return "工具 " + toolName + " 的 MCP 客户端不可用";
        }
        
        // 调用工具并返回结果
        return client.callTool(new McpSchema.CallToolRequest(tool.getToolNameOnServer(), arguments))
                .map(result -> {
                    // result.content() 通常是 List<McpSchema.Content>
                    return result.content().stream()
                            .map(McpSchema.Content::type) // 如果有 text() 方法
                            .filter(Objects::nonNull)
                            .collect(Collectors.joining("\n"));
                })
                .timeout(TIMEOUT)
                .block(); // 阻塞获取 String
    }

    /**
     * 获取所有已暴露工具的函数模式
     * @return 函数模式列表
     */
    public List<FunctionSchema> getFunctionSchemas() {
        // 获取所有已暴露的工具并转换为函数模式列表
        return toolRepository.findAllExposed().stream()
                .map(tool -> new FunctionSchema(
                        tool.getName(),
                        tool.getDescription(),
                        tool.getParamSchema()
                ))
                .collect(Collectors.toList());
    }

    /**
     * 添加或更新 MCP 服务器
     * @param server 要添加或更新的服务器
     * @return 已添加或更新的服务器
     */
    public McpServer addOrUpdateServer(McpServer server) {
        // 保存服务器信息
        McpServer savedServer = serverRepository.save(server);
        // 根据服务器是否启用决定是刷新还是关闭客户端
        if (server.isEnabled()) {
            refreshClient(server.getId());
        } else {
            closeClient(server.getId());
        }
        return savedServer;
    }

    /**
     * 删除 MCP 服务器
     * @param serverId 服务器ID
     */
    public void deleteServer(String serverId) {
        // 关闭客户端
        closeClient(serverId);
        // 删除该服务器的所有工具
        toolRepository.deleteAllByServerId(serverId);
        // 删除服务器本身
        serverRepository.deleteById(serverId);
    }

    /**
     * 获取所有 MCP 服务器
     * @return 所有服务器集合
     */
    public Collection<McpServer> getAllServers() {
        return serverRepository.findAll();
    }

    /**
     * 获取所有 MCP 工具
     * @return 所有工具集合
     */
    public Collection<McpTool> getAllTools() {
        return toolRepository.findAll();
    }

    /**
     * 定时任务：刷新所有 MCP 客户端
     */
    @Scheduled(fixedRate = 3600000) // 每小时执行一次
    public void scheduledRefresh() {
        logger.info("正在执行 MCP 客户端的定时刷新任务");
        refreshAllClients();
    }
}